package nachos.vm;

import java.util.Arrays;

import nachos.machine.*;
import nachos.threads.*;
import nachos.userprog.*;
import nachos.vm.*;

/**
 * A <tt>UserProcess</tt> that supports demand-paging.
 */
public class VMProcess extends UserProcess {
	
	private ExtendedTransEntry[] extendedPageTable;
	/**
	 * 
	 * @param vpn should be valid
	 * @return
	 */
	public ExtendedTransEntry getETE(int vpn) {
		validateVPN(vpn);
		return extendedPageTable[vpn];
	}
	private int currentPageTableSize = 0;
	
    /**
     * Allocate a new process.
     */
    public VMProcess() {
	super();
	tlbSize = Machine.processor().getTLBSize();
    }

    /**
     * Save the state of this process in preparation for a context switch.
     * Called by <tt>UThread.saveState()</tt>.
     */
    public void saveState() {
	super.saveState();
	
		// Sync all tlb entries with process's page table(write tlb back to page table)
		// And invalidate all tlb entries when leaving a process.
		for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			if (tlbEntry.valid) {
				TranslationEntry pageTablEntry = extendedPageTable[tlbEntry.vpn].translationEntry;
				if (tlbEntry.used) pageTablEntry.used = true;
				if (tlbEntry.dirty) pageTablEntry.dirty = true;
				// after sync, invalidate tlb entry
				tlbEntry.valid = false;
				Machine.processor().writeTLBEntry(i, tlbEntry);
			}
		}
    }

    /**
     * Restore the state of this process after a context switch. Called by
     * <tt>UThread.restoreState()</tt>.
     */
    public void restoreState() {

    }

    /**
     * Initializes page tables for this process so that the executable can be
     * demand-paged.
     *
     * @return	<tt>true</tt> if successful.
     */
    protected boolean loadSections() {
    	
//    	if (numPages > Machine.processor().getNumPhysPages()) {
//    	    coff.close();
//    	    Lib.debug(dbgProcess, "\tinsufficient physical memory");
//    	    return false;
//    	}

    	extendedPageTable = new ExtendedTransEntry[numPages];
    	
    	// load sections
    	for (int s=0; s<coff.getNumSections(); s++) {
    	    CoffSection section = coff.getSection(s);
    	    Lib.debug(dbgProcess, "\tgive vpn to " + section.getName()
    			      + " section (" + section.getLength() + " pages) "+this);
    	    for (int i=0; i<section.getLength(); i++) {
    	    	// lazy load, simply add a invalid trans entry
    	    	TranslationEntry te = new TranslationEntry();
    	    	te.readOnly = section.isReadOnly();
    	    	extendedPageTable[currentPageTableSize++] = new ExtendedTransEntry(te);
    	    }
    	}
    	return true;
    }
    
    protected boolean loadStackAndArg() {
    	// allocate pmem for stack
    	for (int i = coffPageNum ; i < coffPageNum + stackPages ; i++) {
    		extendedPageTable[currentPageTableSize++] = new ExtendedTransEntry(new TranslationEntry());
    	}
    	// for arg
    	extendedPageTable[currentPageTableSize++] = new ExtendedTransEntry(new TranslationEntry());
    	return true;
    }
    
    private String[] args;
    private byte[][] argv;
    protected void storeArg(String[] args, byte[][] argv) {
    	this.args = args.clone();
    	this.argv = argv.clone();
    }

    /**
     * Release any resources allocated by <tt>loadSections()</tt>.
     */
    protected void unloadSections() {
	super.unloadSections();
    }    

    /**
     * Handle a user exception. Called by
     * <tt>UserKernel.exceptionHandler()</tt>. The
     * <i>cause</i> argument identifies which exception occurred; see the
     * <tt>Processor.exceptionZZZ</tt> constants.
     *
     * @param	cause	the user exception that occurred.
     */
    public void handleException(int cause) {
	Processor processor = Machine.processor();

	switch (cause) {
	case Processor.exceptionTLBMiss:
		handleTLBMiss();
		break;
	default:
	    super.handleException(cause);
	    break;
	}
    }
    
    private void handleTLBMiss() {
    	int vaddr = Machine.processor().readRegister(Processor.regBadVAddr);
    	int vpn = Processor.pageFromAddress(vaddr);
    	Lib.debug(dbgTLB, "Handle TLB miss for vpn: "+vpn);
    	printTLB();
    	if (vpn < 0 || vpn > numPages - 1) {
    		// illegal vpn! just kill the process
    		Lib.debug(dbgProcess, "illegal vpn:"+vpn+" in handleTLBMiss, kill the process:"+this.toString());
//    		Machine.interrupt().restore(intStatus);
    		handleExit(EXIT_ABNORMALLY);
    		return;
    	}
//    	VMKernel.mainmemLock.acquire();
    	boolean intStatus = Machine.interrupt().disable();
    	ExtendedTransEntry ete = extendedPageTable[vpn];
    	if (ete.translationEntry.valid) {
    		// if te is valid, which means that page is in main mem, simply write te to TLB
    		populateTLB(ete.translationEntry);
    		Machine.interrupt().restore(intStatus);
    	}
    	else {
    		Machine.interrupt().restore(intStatus);
    		// te is invalid, which means that page is **not** in main mem,
    		// load page from swap or Coff section or just blank
    		if (lastProcess != this) { // magic here!!!
    			int wait = 30000;//VMKernel.randomWait.nextInt(5000);
        		ThreadedKernel.alarm.waitUntil(wait);
        		lastProcess = this;
    		}
    		VMKernel.mainmemLock.acquire();
    		loadPage(vpn);
    		populateTLB(ete.translationEntry);
    		VMKernel.mainmemLock.release();
    	}
//    	VMKernel.mainmemLock.release();
    	printTLB();
    	return;
    }
    
    private void printTLB() {
    	Lib.debug(dbgTLB, "<><><><><><><> TLB <><><><><><><>");
    	for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			Lib.debug(dbgTLB, "entry: "+i+" valid: "+tlbEntry.valid+" vpn: "+tlbEntry.vpn+" ppn: "+tlbEntry.ppn+" readonly: "+tlbEntry.readOnly+" used: "+tlbEntry.used+" dirty: "+tlbEntry.dirty);
		}
    	Lib.debug(dbgTLB, "<><><><><><><> END <><><><><><><>");
    }
    
    /**
     * 
     * @param vpn should be valid
     */
    private void loadPage(int vpn) {
    	validateVPN(vpn);
    	if (vpn < coffPageNum) {	// coff page
    		loadOneCoffPage(vpn);
    	} else if (vpn == numPages - 1) { // arg page
    		loadArgPage(vpn);
    	} else {					// stack page
    		loadOneStackPage(vpn);
    	}
    }
    
    private void loadOneCoffPage(int vpn) {
    	// find the coff page
    	int count = 0;
    	for (int s=0; s<coff.getNumSections(); s++) {
    	    CoffSection section = coff.getSection(s);

    	    for (int i=0; i<section.getLength(); i++) {
    		
    	    	if (count++ != vpn)
    	    		continue;
    	    	
    	    	// have found the coff page, load it to mem
    	    	
    	    	// get a free page
    	    	int ppn = VMKernel.getPhyPage(this, vpn);
    	    	if (ppn == -1) {
        	    	Lib.debug(dbgProcess, "loadOneCoffPage can't get phy page");
        	    	return;
        	    }
    	    	
        	    Lib.debug(dbgProcess, "\tinitializing " + section.getName()
        	    		+ " section (page: "+i+") "+" get ppn: "+ppn+" "+this);
    	    	
    	    	// init page with data
    	    	ExtendedTransEntry ete = extendedPageTable[vpn];
    	    	if (ete.translationEntry.readOnly) {
    	    		// if read only, surely not in swap, simply load from section file
    	    		section.loadPage(i, ppn); 
    	    	} else {
    	    		// if not read only, some data many have changed before and moved to swap
    	    		// so check if swap exists
    	    		if (ete.spn == -1) {
    	    			// not in swap
    	    			section.loadPage(i, ppn);
    	    		} else {
    	    			// in swap
    	    			VMKernel.swap.restoreFromSwap(ppn, ete.spn);
    	    		}
    	    	}

    	    	// create trans entry mapping
    	    	refreshETE(ete, vpn, ppn);
    	    	
    	    	// populate TLB!
//    	    	populateTLB(ete.translationEntry);
    	    }
    	}
    }
	
    private void loadOneStackPage(int vpn) {
    	// get a free page
    	int ppn = VMKernel.getPhyPage(this, vpn);
    	if (ppn == -1) {
	    	Lib.debug(dbgProcess, "loadOneStackPage can't get phy page");
	    	return;
	    }
    	Lib.debug(dbgProcess, "\tloadOneStackPage: " + vpn + " get ppn:"+ppn+ " "+this);
    	// init page with data
    	ExtendedTransEntry ete = extendedPageTable[vpn]; 
    	
    	// some data many have changed before and moved to swap
		// so check if swap exists
		if (ete.spn == -1) {
			// not in swap, fill with zero
			Arrays.fill(Machine.processor().getMemory(), ppn*pageSize, (ppn+1)*pageSize, (byte) 0);
		} else {
			// in swap
			VMKernel.swap.restoreFromSwap(ppn, ete.spn);
		}

    	// create trans entry mapping
    	refreshETE(ete, vpn, ppn);
    	
    	// populate TLB!
//    	populateTLB(ete.translationEntry);
    }
    
    private void loadArgPage(int vpn) {
    	
    	// get a free page
    	int ppn = VMKernel.getPhyPage(this, vpn);
    	if (ppn == -1) {
	    	Lib.debug(dbgProcess, "loadArgPage can't get phy page");
	    	return;
	    }
    	
    	// init page with data
    	ExtendedTransEntry ete = extendedPageTable[vpn]; 
    	
    	// create trans entry mapping
    	refreshETE(ete, vpn, ppn); // create mapping first for storeArg
    	
    	// some data many have changed before and moved to swap
		// so check if swap exists
		if (ete.spn == -1) {
			// not in swap
			super.storeArg(args, argv);
			refreshETE(ete, vpn, ppn); // clear used and dirty bits that is changed by storeArg
		} else {
			// in swap
			VMKernel.swap.restoreFromSwap(ppn, ete.spn);
		}
		
    	// populate TLB!
//    	populateTLB(ete.translationEntry);
    	
    }
    
    private void refreshETE(ExtendedTransEntry ete, int vpn, int ppn) {
    	ete.translationEntry.vpn = vpn;
    	ete.translationEntry.ppn = ppn;
    	ete.translationEntry.valid = true;
    	ete.translationEntry.used = true;
    	ete.translationEntry.dirty = false;
    }
    
    /**
     * randomly pick up a TLB entry from processor, write the entry(dirty and used bit) back to process's page table
     * replace that entry with the param
     * @param translationEntry
     */
    private void populateTLB(TranslationEntry translationEntry) {
    	int tlbWritePos = VMKernel.getTLBWritePos(); // random
    	// write back the old entry
    	TranslationEntry oldEntryInTLB = Machine.processor().readTLBEntry(tlbWritePos);
    	if (oldEntryInTLB.valid) {
    		TranslationEntry oldEntryInPageTable = extendedPageTable[oldEntryInTLB.vpn].translationEntry;
        	oldEntryInPageTable.dirty = oldEntryInTLB.dirty;
        	oldEntryInPageTable.used = oldEntryInTLB.used;
    	}
    	
    	// replace with new entry
    	Machine.processor().writeTLBEntry(tlbWritePos, translationEntry);
    }
    
    /**
     * 
     * @param vaddr
     * @param isWrite
     * @return the paddr, -1 if write to read only page, or no pageTable entry for the vaddr 
     */
    protected int translate(int vaddr, boolean isWrite) {
    	int vpn = Processor.pageFromAddress(vaddr);
    	
    	if (vpn < 0 || vpn > numPages - 1) {
    		// illegal vpn!
    		return -1;
    	}
    	
    	TranslationEntry entry = extendedPageTable[vpn].translationEntry;
    	
    	if (entry.readOnly && isWrite) {
    		return -1;
    	}
    	int ppn = entry.ppn;
    	int offset = Processor.offsetFromAddress(vaddr);
    	return Processor.makeAddress(ppn, offset);
    }
    
    /**
     * if the page is not currently in main mem, initialize it or get it from swap
     * @param vpn should be valid
     */
    protected void pinPage(int vpn) {
    	validateVPN(vpn);
    	boolean hasLock = VMKernel.mainmemLock.isHeldByCurrentThread();
    	if (!hasLock) {
    		VMKernel.mainmemLock.acquire();
    	}
    	
    	extendedPageTable[vpn].isPinned = true;
    	
    	//  check if the page is in main mem
    	if (!extendedPageTable[vpn].translationEntry.valid) {
    		// if not, load page in
    		loadPage(vpn);
    	}
    	if (!hasLock) {
    		VMKernel.mainmemLock.release();
    	}
    }
    
    /**
     * also set used/dirty bits properly in page table and sync with TLB if possible
     * @param vpn should be valid
     * @param isWrite if true, set dirty bit in entry
     */
    protected void unPinPage(int vpn, boolean isWrite) {
    	validateVPN(vpn);
    	boolean hasLock = VMKernel.mainmemLock.isHeldByCurrentThread();
    	if (!hasLock) {
    		VMKernel.mainmemLock.acquire();
    	}
    	
    	TranslationEntry pageTableEntry = extendedPageTable[vpn].translationEntry;
    	pageTableEntry.used = true;
    	if (isWrite) {
    		pageTableEntry.dirty = true;
    	}
    	// sync with TLB if possible
		for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			if (tlbEntry.vpn == vpn && tlbEntry.valid) {
				if (tlbEntry.dirty)
					pageTableEntry.dirty = true;
				Machine.processor().writeTLBEntry(i, pageTableEntry);
			}
		}
		
    	extendedPageTable[vpn].isPinned = false;
    	if (!hasLock) {
    		VMKernel.mainmemLock.release();
    	}
    	// signal possible waiter on pinned page
    	VMKernel.pinWaiter.releasePin();
    }
    
    public void evictPage(int vpn) {
    	validateVPN(vpn);
    	ExtendedTransEntry ete = extendedPageTable[vpn];
    	// check dirty bit from TLB again
    	boolean intStatus = Machine.interrupt().disable();
    	for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			if (tlbEntry.valid && tlbEntry.vpn==vpn) {
				if (tlbEntry.dirty) {
					ete.translationEntry.dirty = true;
				}
			}
		}    	
    	
    	if (!ete.translationEntry.readOnly && ete.translationEntry.dirty) {
    		// not read only and is dirty, write to swap
    		ete.spn = VMKernel.swap.saveToSwap(ete.translationEntry.ppn, ete.spn);
    	}
    	Machine.interrupt().restore(intStatus);
    	ete.translationEntry.valid = false;
    	// should also invalidate corresponding TLB entry!!!
    	for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			if (tlbEntry.valid && tlbEntry.vpn == vpn) {
				tlbEntry.valid = false;
				Machine.processor().writeTLBEntry(i, tlbEntry);
			}
		}
    }
    
    private void validateVPN(int vpn) {
    	if (vpn<0 || vpn>=numPages) {
    		handleExit(EXIT_ABNORMALLY);
//			Lib.assertNotReached("Invalid VPN");
		}
    }
    
    protected void freeAllMemory() {
    	// invalidate all TLB !!!
    	for (int i=0 ; i<tlbSize ; i++) {
			TranslationEntry tlbEntry = Machine.processor().readTLBEntry(i);
			if (tlbEntry.valid) {
				tlbEntry.valid = false;
				Machine.processor().writeTLBEntry(i, tlbEntry);
			}
		}
    	
    	if (extendedPageTable == null) {
    		return;
    	}
    	VMKernel.mainmemLock.acquire();
    	for (int i = 0 ; i < currentPageTableSize ; i++) {
    		TranslationEntry translationEntry = extendedPageTable[i].translationEntry;
    		if (translationEntry.valid) {
    			VMKernel.freePhyPage(translationEntry.ppn);
    		} 
    		// free swap page
    		VMKernel.swap.freeSwapPage(extendedPageTable[i].spn);
    	}
    	VMKernel.mainmemLock.release();
    }
    
    private int tlbSize;
    private static final int pageSize = Processor.pageSize;
    private static final char dbgProcess = 'a';
    private static final char dbgVM = 'v';
    private static final char dbgTLB = 'l';
    
    private static VMProcess lastProcess = null;
}
